#include "parser.h"
#include <cassert>

/*
prog : stmt*
stmt : decl-stmt | expr-stmt | null-stmt | if-stmt
null-stmt : ";"
decl-stmt : "int" identifier ("," identifier (= expr)?)* ";"
if-stmt : "if" "(" expr ")" stmt ( "else" stmt )?
expr-stmt : expr ";"
expr : assign-expr | add-expr
assign-expr: identifier "=" expr
add-expr : mult-expr (("+" | "-") mult-expr)*
mult-expr : primary-expr (("*" | "/") primary-expr)*
primary-expr : identifier | number | "(" expr ")"
number: ([0-9])+
identifier : (a-zA-Z_)(a-zA-Z0-9_)*
*/

// 解析目标程序
// stmt : decl-stmt | expr-stmt | null-stmt
std::shared_ptr<Program> Parser::ParseProgram()
{
    std::vector<std::shared_ptr<ASTNode>> nodeVec;
    while (token.tokenType != TokenType::eof)
    {
        auto stmt = ParseStmt();
        if (stmt)
            nodeVec.push_back(stmt);
    }
    auto program = std::make_shared<Program>();
    program->nodeVec = std::move(nodeVec);
    return program;
}

// 解析语句
std::shared_ptr<ASTNode> Parser::ParseStmt()
{
    // 遇到 ; 需要进行消费 token
    // null-stmt
    if (token.tokenType == TokenType::semi)
    {
        Consume(TokenType::semi);
        return nullptr;
    }
    // decl-stmt
    else if (IsTypeName())
    {
        return ParseDeclStmt();
    }
    // block-stmt
    else if (token.tokenType == TokenType::l_brace)
    {
        return ParseBlockStmt();
    }
    // if-stmt
    else if (token.tokenType == TokenType::kw_if)
    {
        return ParseIfStmt();
    }
    // for-stmt
    else if (token.tokenType == TokenType::kw_for)
    {
        return ParseForStmt();
    }
    // break-stmt
    else if (token.tokenType == TokenType::kw_break)
    {
        return ParseBreakStmt();
    }
    // continue-stmt
    else if (token.tokenType == TokenType::kw_continue)
    {
        return ParseContinueStmt();
    }
    // expr-stmt
    else
    {
        return ParseExprStmt();
    }
}

// 解析声明语句
std::shared_ptr<ASTNode> Parser::ParseDeclStmt()
{
    /// int a, b = 3;
    /// int a = 3;
    Consume(TokenType::kw_int);
    CType *baseTy = CType::getIntTy();
    /// a , b = 3;
    /// a = 3;

    auto declStmt = std::make_shared<DeclStmt>();

    /// a, b = 3;
    /// a = 3;
    int i = 0;
    while (token.tokenType != TokenType::semi)
    {
        if (i++ > 0) // if (i++)
        {
            assert(Consume(TokenType::comma));
        }

        /// 变量声明的节点: int a = 3; -> int a; a = 3;
        // a = 3;
        auto variableDecl = sema.SemaVariableDeclNode(token, baseTy); // get a type
        declStmt->nodeVec.push_back(variableDecl);

        Token tmp = token;
        Consume(TokenType::indentifier);

        // = 3;
        if (token.tokenType == TokenType::equal)
        {
            Token opToken = token;
            Advance();

            // 3;
            auto right = ParseExpr();
            auto left = sema.SemaVariableAccessNode(tmp);
            auto assign = sema.SemaAssignExprNode(left, right, opToken);

            declStmt->nodeVec.push_back(assign);
        }
    }

    Consume(TokenType::semi);

    return declStmt;
}

std::shared_ptr<ASTNode> Parser::ParseBlockStmt()
{
    sema.EnterScope(); // 进入作用域

    auto blockStmt = std::make_shared<BlockStmt>();
    Consume(TokenType::l_brace);
    while (token.tokenType != TokenType::r_brace)
    {
        auto stmt = ParseStmt();
        if (stmt)
            blockStmt->nodeVec.push_back(stmt);
    }
    Consume(TokenType::r_brace);

    sema.ExitScope(); // 离开作用域

    return blockStmt;
}

// 解析表达式语句
std::shared_ptr<ASTNode> Parser::ParseExprStmt()
{
    auto expr = ParseExpr();
    Consume(TokenType::semi);
    return expr;
}

// if-stmt : "if" "(" expr ")" stmt ( "else" stmt )?
/*
if (a)
  b = 3;
else
    b = 4;
*/
std::shared_ptr<ASTNode> Parser::ParseIfStmt()
{
    Consume(TokenType::kw_if);
    Consume(TokenType::l_parent);
    auto condExpr = ParseExpr();
    Consume(TokenType::r_parent);
    auto thenStmt = ParseStmt();
    std::shared_ptr<ASTNode> elseStmt = nullptr;
    if (token.tokenType == TokenType::kw_else)
    {
        Consume(TokenType::kw_else);
        elseStmt = ParseStmt();
    }
    return sema.SemaIfStmtNode(condExpr, thenStmt, elseStmt);
}

/*
for-stmt :
"for" "(" expr? ";" expr? ";" expr? ")" stmt
"for" "(" decl-stmt expr? ";" expr? ")" stmt
*/
// 解析 for-stmt
std::shared_ptr<ASTNode> Parser::ParseForStmt()
{
    sema.EnterScope();

    Consume(TokenType::kw_for);
    Consume(TokenType::l_parent);

    auto forStmt = std::make_shared<ForStmt>();

    breakNodes.push_back(forStmt);
    continueNodes.push_back(forStmt);

    std::shared_ptr<ASTNode> initNode = nullptr;
    std::shared_ptr<ASTNode> condNode = nullptr;
    std::shared_ptr<ASTNode> incNode = nullptr;
    std::shared_ptr<ASTNode> bodyNode = nullptr;

    if (IsTypeName())
    {
        initNode = ParseDeclStmt();
    }
    else
    {
        initNode = ParseExpr();
        Consume(TokenType::semi);
    }

    condNode = ParseExpr();
    Consume(TokenType::semi);

    incNode = ParseExpr();
    Consume(TokenType::r_parent);

    bodyNode = ParseStmt();

    forStmt->initNode = initNode;
    forStmt->condNode = condNode;
    forStmt->incNode = incNode;
    forStmt->bodyNode = bodyNode;

    breakNodes.pop_back();
    continueNodes.pop_back();

    sema.ExitScope();

    return forStmt;
}

std::shared_ptr<ASTNode> Parser::ParseBreakStmt()
{
    if (breakNodes.size() == 0)
    {
        GetDiagEngine().Report(llvm::SMLoc::getFromPointer(token.ptr), diag::err_break_stmt);
    }
    Consume(TokenType::kw_break);
    auto node = std::make_shared<BreakStmt>();
    node->target = breakNodes.back();
    return node;
}

std::shared_ptr<ASTNode> Parser::ParseContinueStmt()
{
    if (continueNodes.size() == 0)
    {
        GetDiagEngine().Report(llvm::SMLoc::getFromPointer(token.ptr), diag::err_continue_stmt);
    }
    Consume(TokenType::kw_continue);
    auto node = std::make_shared<ContinueStmt>();
    node->target = continueNodes.back();
    return node;
}

// 解析表达式
// expr : assign-expr | add-expr
// assign-expr: identifier "=" expr
// add-expr : mult-expr (("+" | "-") mult-expr)*
std::shared_ptr<ASTNode> Parser::ParseExpr()
{
    lexer.SaveState();
    bool isAssign = false;
    Token tmp = token;
    // a = b;
    if (tmp.tokenType == TokenType::indentifier)
    {
        lexer.NextToken(tmp);
        if (tmp.tokenType == TokenType::equal)
        {
            isAssign = true;
        }
    }
    lexer.RestoreState();
    if (isAssign)
    {
        return ParseAssignExpr();
    }
    else // logor-expr
        return ParseLogOrExpr();
}

// 解析赋值表达式
std::shared_ptr<ASTNode> Parser::ParseAssignExpr()
{
    // a = b;
    Token tmp = token;
    Consume(TokenType::indentifier);
    auto expr = sema.SemaVariableAccessNode(tmp);
    Token opToken = token;
    Consume(TokenType::equal);
    return sema.SemaAssignExprNode(expr, ParseExpr(), opToken);
}

// 解析逻辑或表达式 ||
std::shared_ptr<ASTNode> Parser::ParseLogOrExpr()
{
    std::shared_ptr<ASTNode> left = ParseLogAndExpr();
    while (token.tokenType == TokenType::pipepipe)
    {
        OpCode op = OpCode::log_or;
        Advance();
        auto right = ParseLogAndExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析逻辑与表达式 &&
std::shared_ptr<ASTNode> Parser::ParseLogAndExpr()
{
    std::shared_ptr<ASTNode> left = ParseBitOrExpr();
    while (token.tokenType == TokenType::ampamp)
    {
        OpCode op = OpCode::log_and;
        Advance();
        auto right = ParseBitOrExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析按位或表达式 |
std::shared_ptr<ASTNode> Parser::ParseBitOrExpr()
{
    std::shared_ptr<ASTNode> left = ParseBitXorExpr();
    while (token.tokenType == TokenType::pipe)
    {
        OpCode op = OpCode::bit_or;
        Advance();
        auto right = ParseBitXorExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析按位异或表达式 ^
std::shared_ptr<ASTNode> Parser::ParseBitXorExpr()
{
    std::shared_ptr<ASTNode> left = ParseBitAndExpr();
    while (token.tokenType == TokenType::caret)
    {
        OpCode op = OpCode::bit_xor;
        Advance();
        auto right = ParseBitAndExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析按位与表达式 &
std::shared_ptr<ASTNode> Parser::ParseBitAndExpr()
{
    std::shared_ptr<ASTNode> left = ParseEqualExpr();
    while (token.tokenType == TokenType::amp)
    {
        OpCode op = OpCode::bit_and;
        Advance();
        auto right = ParseEqualExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析等号表达式
std::shared_ptr<ASTNode> Parser::ParseEqualExpr()
{
    std::shared_ptr<ASTNode> left = ParseRelationalExpr();
    while (token.tokenType == TokenType::equal_equal || token.tokenType == TokenType::not_equal)
    {
        OpCode op;
        if (token.tokenType == TokenType::equal_equal)
        {
            op = OpCode::equal_equal;
        }
        else
        {
            op = OpCode::not_equal;
        }
        Advance();
        auto right = ParseRelationalExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析关系表达式
std::shared_ptr<ASTNode> Parser::ParseRelationalExpr()
{
    std::shared_ptr<ASTNode> left = ParseShiftExpr();
    while (token.tokenType == TokenType::less || token.tokenType == TokenType::less_equal || token.tokenType == TokenType::greater || token.tokenType == TokenType::greater_equal)
    {
        OpCode op;
        if (token.tokenType == TokenType::less)
        {
            op = OpCode::less;
        }
        else if (token.tokenType == TokenType::less_equal)
        {
            op = OpCode::less_equal;
        }
        else if (token.tokenType == TokenType::greater)
        {
            op = OpCode::greater;
        }
        else
        {
            op = OpCode::greater_equal;
        }
        Advance();
        auto right = ParseShiftExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

std::shared_ptr<ASTNode> Parser::ParseShiftExpr()
{
    std::shared_ptr<ASTNode> left = ParseAddExpr();
    while (token.tokenType == TokenType::less_less || token.tokenType == TokenType::greater_greater)
    {
        OpCode op;
        if (token.tokenType == TokenType::less_less) // <<
        {
            op = OpCode::left_shift;
        }
        else // >>
        {
            op = OpCode::right_shift;
        }
        Advance();
        auto right = ParseAddExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}


// 解析加法表达式
std::shared_ptr<ASTNode> Parser::ParseAddExpr()
{
    std::shared_ptr<ASTNode> left = ParseMultiExpr();
    while (token.tokenType == TokenType::plus || token.tokenType == TokenType::minus)
    {
        OpCode op;
        if (token.tokenType == TokenType::plus)
        {
            op = OpCode::add;
        }
        else
        {
            op = OpCode::sub;
        }
        Advance();
        auto right = ParseMultiExpr();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析项
std::shared_ptr<ASTNode> Parser::ParseMultiExpr()
{
    std::shared_ptr<ASTNode> left = ParsePrimary();

    while (token.tokenType == TokenType::star || token.tokenType == TokenType::slash || token.tokenType == TokenType::percent)
    {
        OpCode op;
        if (token.tokenType == TokenType::star)
        {
            op = OpCode::mul;
        }
        else if (token.tokenType == TokenType::slash)
        {
            op = OpCode::div;
        }
        else
        {
            op = OpCode::mod;
        }
        Advance();
        auto right = ParsePrimary();
        auto binaryExpr = sema.SemaBinaryExprNode(left, right, op);

        left = binaryExpr;
    }
    return left;
}

// 解析因子
std::shared_ptr<ASTNode> Parser::ParsePrimary()
{
    if (token.tokenType == TokenType::l_parent)
    {
        Consume(TokenType::l_parent);
        auto expr = ParseExpr();
        Consume(TokenType::r_parent);
        return expr;
    }
    else if (token.tokenType == TokenType::indentifier)
    {
        auto variableAccessExpr = sema.SemaVariableAccessNode(token);
        Consume(TokenType::indentifier);
        return variableAccessExpr;
    }
    else
    {
        Expect(TokenType::number);
        auto factorExpr = sema.SemaNumberExprNode(token, token.type);
        Consume(TokenType::number);
        return factorExpr;
    }
}

/// 消耗 token 函数
bool Parser::Expect(TokenType tokenType)
{
    if (token.tokenType == tokenType)
        return true;
    GetDiagEngine().Report(llvm::SMLoc::getFromPointer(token.ptr),
                           diag::err_expected,
                           Token::GetSpellingText(tokenType),
                           llvm::StringRef(token.ptr, token.length));
    return false;
}

bool Parser::Consume(TokenType tokenType)
{
    if (Expect(tokenType))
    {
        Advance();
        return true;
    }
    return false;
}

bool Parser::Advance()
{
    lexer.NextToken(token);
    return true;
}

bool Parser::IsTypeName()
{
    if (token.tokenType == TokenType::kw_int)
    {
        return true;
    }
    return false;
}